#include <gsl/gsl_blas.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_matrix.h>
#include <math.h>
#include <mex.h>
#include <omp.h>
#include <stdio.h>

//[E,z] = (n, n_class, pi, K);
void mexFunction(int nlhs, mxArray* plhs[], int nrhs, const mxArray* prhs[])

{
	
	double sigma_noise = 5e-7;
	double *pi;
	double *Kmx;
	int c, i, n, n_class;
	double *tmp, z = 0;
	
	 /* Check for proper number of arguments. */
	if(nrhs!=4) {
	    mexErrMsgTxt("mex_laplace_subroutine: requires exactly four input arguments.");
	} else if(nlhs>2) {
	    mexErrMsgTxt("mex_laplace_subroutine: requires at most two output arguments.");
	}
	
	tmp = mxGetPr(prhs[0]);
	n = (int) (tmp[0]);
	tmp = mxGetPr(prhs[1]);
	n_class = (int) (tmp[0]);
	pi = mxGetPr(prhs[2]);
	Kmx = mxGetPr(prhs[3]);
	
	gsl_matrix *K = gsl_matrix_alloc(n*n_class, n*n_class);
	K->data = Kmx;
	
	gsl_matrix *E = gsl_matrix_alloc(n*n_class, n);
	
	gsl_matrix* Ip = gsl_matrix_alloc(n, n);
	gsl_matrix_set_identity(Ip);
	gsl_matrix_scale(Ip, 1+sigma_noise);
	
	//#pragma omp parallel for
	for (c = 0; c < n_class; c++) {
	    gsl_matrix *sqrtDc = gsl_matrix_calloc(n, n);
	    gsl_matrix *L = gsl_matrix_alloc(n, n);
	    gsl_matrix *L1 = gsl_matrix_alloc(n, n);
	    _gsl_matrix_const_view Kcv = gsl_matrix_const_submatrix(K, c*n, c*n, n, n);
	    gsl_matrix *Kc = gsl_matrix_alloc(n, n);
	    Kc = &Kcv.matrix;
	    //printf("Kc's dims are %d %d\n", Kc->size1, Kc->size2);
	    _gsl_matrix_view Ecv = gsl_matrix_submatrix(E, c*n, 0, n, n);
	    gsl_matrix *Ec = &Ecv.matrix;
	    for (i = 0; i < n; i++) {
		gsl_matrix_set(sqrtDc, i, i, sqrt(pi[i + c * n]));
	    }
	    
	    gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, Kc, sqrtDc, 0.0, L1); // L = K[:,:,c]) * sqrtDc 
	    gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, sqrtDc, L1, 0.0, L); // L = sqrtDc * K(:,:,c) * sqrtDc;
	    // NTS: is the above equivalent to:
	    // Lc := (pi_c pi_c') .* Kc
	    // ?
	    gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, Ip, Ip, 1.0, L); // L = (1 + sigma_noise)*eye(n) + sqrtDc * K(:,:,c) * sqrtDc;
	    // rewrite that one, god - mult to do add? that's just absurd

	    gsl_linalg_cholesky_decomp(L); // L = chol((1 + sigma_noise)*eye(n) + sqrtDc * K(:,:,c) * sqrtDc)'; 
	    // this gives L as being lower triangular, ie L'L = RHS
	    // if it throws GSL_EDOM, that means L wasn't pos def, apparently
	    
	    /*
	    if (c == 0) {
	    int j;
	    for (i = 0; i < n; i++) {
		for (j = 0; j < n; j++) {
		    printf("%g ", gsl_matrix_get(L, i, j));
		}
		printf("\n");
	    }*/

	    for (i = 0; i < n; i++) {
		z += log(gsl_matrix_get(L, i, i));
	    }
	    
	    Ec = sqrtDc; // are you sure this is a deep copy?
	    
	    gsl_blas_dtrsm(CblasLeft, CblasLower, CblasNoTrans, CblasUnit, 1.0, L, Ec); // Ec = L \ sqrtDc
	    gsl_blas_dtrsm(CblasLeft, CblasLower, CblasTrans, CblasUnit, 1.0, L, Ec); // Ec = L' \ (L \ sqrtDc) 
	    gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, sqrtDc, Ec, 0.0, Ec); // Ec = sqrtDc * ( L' \ (L \ sqrtDc) );
	    
	    /*catch
		disp('alg3.3s chol is failing again');
		matrix [L, U] = lu((1 + sigma_noise)*eye(n) + sqrtDc * K(:,:,c) * sqrtDc);
		E(:,:,c) = ((sqrtDc / U) / L) * sqrtDc;*/
	    //end
	}
	if (nlhs >= 1) {
	    mxArray *v = mxCreateDoubleMatrix (n*n_class, n, mxREAL);
	    double *data = mxGetPr(v);
	    data = E->data;
	    plhs[0] = v;
	}
	if (nlhs == 2) {
	    mxArray *zp = mxCreateDoubleMatrix(1, 1, mxREAL);
	    double *zdata = mxGetPr(zp);
	    zdata = &z;
	    plhs[1] = zp;
	}
 
}


// and we're golden. compile with 
// mkoctfile --mex -Wall -I/usr/local/include -L/usr/lib -lgsl -lgslcblas mex_laplace_subroutine.c
